# -*- coding: utf-8 -*-
import time, datetime
import pysnowball as ball
import pandas as pd
from pysnowball import api_ref
from pysnowball import utls

from infrastructure.utils.enums import Period, StockTradeDataName, DateFormat
from config.baseConfiguration import XueQiuConfig

ball.set_token(f"xq_a_token={XueQiuConfig.xueqiu_token};u={XueQiuConfig.u}")

delta8hours = datetime.timedelta(hours=8)


def kline(symbol: str, period : Period = Period.day, days=5) -> pd.DataFrame:
    time.sleep(2)
    code = symbol
    beginTime = int(time.time() * 1000 + 24 * 60 * 60 * 1000)
    multiple = 1
    if period == Period.minute30:
        multiple = 4 * 2
    elif period == Period.minute5:
        multiple = 4 * 2 * 6
    elif period == Period.minute1:
        multiple = 4 * 2 * 6 * 1
    count = days*multiple
    url = api_ref.kline.format(code, beginTime, count)
    if period == Period.minute30:
        url = url.replace("period=day", "period=30m")
    elif period == Period.minute5:
        url = url.replace("period=day", "period=5m")
    elif period == Period.minute1:
        url = url.replace("period=day", "period=1m")
    try:
        result = utls.fetch(url)
    except Exception as e:
        print(e)
        time.sleep(2)
        result = utls.fetch(url)
    df = pd.DataFrame(result["data"]["item"], columns=result["data"]["column"])
    # logger.info(f"下载股票{symbol} {period.name}周期数据：{df.shape[0]}条")
    df = df.drop([
        "volume_post",
        "amount_post",
        "pe",
        "pb",
        "ps",
        "pcf",
        "market_capital",
        "balance",
        "hold_volume_cn",
        "hold_ratio_cn",
        "net_volume_cn",
        "hold_volume_hk",
        "hold_ratio_hk",
        "net_volume_hk"
             ], axis=1)
    df[StockTradeDataName.DATE_INDEX.value] = df["timestamp"].apply(lambda x: datetime.datetime.utcfromtimestamp(x / 1000) + delta8hours)
    df[StockTradeDataName.TRADE_DATE.value] = df[StockTradeDataName.DATE_INDEX.value].apply(lambda x: int(x.strftime(DateFormat.YmdHMS.value)))
    df = df.drop("timestamp", axis=1)
    df.set_index(keys=StockTradeDataName.DATE_INDEX.value, inplace=True)
    df = df.rename(columns={
        "volume": StockTradeDataName.VOL.value,
        "open": StockTradeDataName.OPEN.value,
        "high": StockTradeDataName.HIGH.value,
        "low": StockTradeDataName.LOW.value,
        "close": StockTradeDataName.CLOSE.value,
        "chg": StockTradeDataName.CHANGE.value,
        "percent": StockTradeDataName.PCT_CHG.value,
        "turnoverrate": StockTradeDataName.TURNOVER_RATE.value,
        "amount": StockTradeDataName.AMOUNT.value
    })
    return df

if __name__ == '__main__':
    df = kline("SZ000001", period=Period.minute5)
    print(df)